library(ggplot2)
library(reshape2)
library(latex2exp)

save_path <- "../../paper/"
height    <- 5
width     <- 5

x <- as.matrix(seq(-5, 5, 0.1))

hinge_loss <- function(x) {
  y <- x
  y[x < 0] <- 0
  return(y)
}

squared_hinge_loss <- function(x) {
  y <- x^2
  y[x < 0] <- 0
  return(y)
}

als_loss <- function(x, p) {
  y <- x^2
  y[x >= 0] <- p*y[x >= 0]
  y[x <  0] <- (1-p)*y[x <  0]
  return(y)
}

bals_loss <- function(x, eta, lambda, p) {
  y <- (1 - 1 / (1 + lambda*als_loss(x, p)))/eta
  return(y)
}

hinge_y <- hinge_loss(x)
squared_hinge_y <- squared_hinge_loss(x)

df_fig_1 <- as.data.frame(cbind(x, hinge_y, squared_hinge_y))
colnames(df_fig_1) <- c("x", "hinge", "sh")

df_fig_1 <- melt(df_fig_1, id.vars = "x")

ggplot(df_fig_1, aes(x = x, y = value, color = variable, linetype = variable)) + 
  geom_line() +
  theme_bw() +
  xlab("z") +
  ylab("Loss value") + 
  scale_color_manual(values = c(4, 2),
                     labels = c("Hinge loss function", 
                                "Squared hinge loss function")) +
  scale_linetype_manual(values = c(4, 5),
                        labels = c("Hinge loss function", 
                                  "Squared hinge loss function")) +
  theme(legend.position = "inside",
        legend.position.inside = c(0.25, 0.20)) +
  guides(color = guide_legend(ncol = 1, byrow = TRUE, title = NULL),
         linetype = guide_legend(ncol = 1, byrow = TRUE, title = NULL)) +
  xlim(-5, 5)
 
fig_path <- paste(save_path, "Hinge_and_Squared_Hinge.eps", sep = "")
ggsave(fig_path, height = height, width = width)


bals1 <- bals_loss(x, 1, 1, 0.9)
bals2 <- bals_loss(x, 2, 1, 0.5)
bals3 <- bals_loss(x, 1, 1, 0.1)
bals4 <- bals_loss(x, 1, 2, 0.9)
bals_df <- as.data.frame(cbind(x, bals1, bals2, bals3, bals4))
colnames(bals_df) <- c("x", "bals1", "bals2", "bals3", "bals4")
bals_df <- melt(bals_df, id.vars = "x")

ggplot(bals_df, aes(x = x, y = value, linetype = variable, color = variable)) + 
  geom_line() +
  theme_bw() +
  xlab("z") +
  ylab("Loss value") + 
  theme(legend.position = "inside",
        legend.position.inside = c(0.20, 0.20)) +
  scale_color_discrete(
                       labels = c(TeX("$\\eta = 1$, $\\lambda = 1$, $p = 0.9$"), 
                                  TeX("$\\eta = 2$, $\\lambda = 1$, $p = 0.5$"),
                                  TeX("$\\eta = 1$, $\\lambda = 1$, $p = 0.1$"),
                                  TeX("$\\eta = 1$, $\\lambda = 2$, $p = 0.9$"))) +
  scale_linetype_discrete(
                          labels = c(TeX("$\\eta = 1$, $\\lambda = 1$, $p = 0.9$"), 
                                     TeX("$\\eta = 2$, $\\lambda = 1$, $p = 0.5$"),
                                     TeX("$\\eta = 1$, $\\lambda = 1$, $p = 0.1$"),
                                     TeX("$\\eta = 1$, $\\lambda = 2$, $p = 0.9$"))) +
  guides(color = guide_legend(ncol = 1, byrow = TRUE, title = NULL),
         linetype = guide_legend(ncol = 1, byrow = TRUE, title = NULL)) +
  xlim(-5, 5)

fig_path <- paste(save_path, "BALS_Loss_Diff_Param.eps", sep = "")
ggsave(fig_path, height = height, width = width)


bsh1 <- bals_loss(x, 1, 0.5, 1)
bsh2 <- bals_loss(x, 2, 1, 1)
bsh3 <- bals_loss(x, 1, 2, 1)
bsh4 <- bals_loss(x, 1, 4, 1)
bals_df <- as.data.frame(cbind(x, bsh1, bsh2, bsh3, bsh4))
colnames(bals_df) <- c("x", "bsh1", "bsh2", "bsh3", "bsh4")
bsh_df <- melt(bals_df, id.vars = "x")

ggplot(bsh_df, aes(x = x, y = value, linetype = variable, color = variable)) + 
  geom_line() +
  theme_bw() +
  xlab("z") +
  ylab("Loss value") + 
  theme(legend.position = "inside",
        legend.position.inside = c(0.20, 0.20)) +
  scale_color_discrete(
    labels = c(TeX("$\\eta = 1$, $\\lambda = 0.5$$"), 
               TeX("$\\eta = 2$, $\\lambda = 1$"),
               TeX("$\\eta = 1$, $\\lambda = 2$"),
               TeX("$\\eta = 1$, $\\lambda = 4$"))) +
  scale_linetype_discrete(
    labels = c(TeX("$\\eta = 1$, $\\lambda = 0.5$$"), 
               TeX("$\\eta = 2$, $\\lambda = 1$"),
               TeX("$\\eta = 1$, $\\lambda = 2$"),
               TeX("$\\eta = 1$, $\\lambda = 4$"))) +
  guides(color = guide_legend(ncol = 1, byrow = TRUE, title = NULL),
         linetype = guide_legend(ncol = 1, byrow = TRUE, title = NULL)) +
  xlim(-5, 5)

fig_path <- paste(save_path, "BSH_Loss_Diff_Param.eps", sep = "")
ggsave(fig_path, height = height, width = width)